##########################################################################################

library(data.table)
library(optparse)
library(parallel)
library(dplyr)
library(sva)
suppressWarnings(library(edgeR, quietly = T))
library(GenomicFeatures)
library(ggplot2)
library("FactoMineR")
library("factoextra")

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--sample_list_public_file"), type = "character"),
    make_option(c("--raw_tcga_count_file"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--gtf_for_len_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    sample_list_file <- "~/20220915_gastric_multiple/rna_combine/analysis/config/tumor_normal.list"
    rsem_file <- "~/20220915_gastric_multiple/rna_combine/analysis/RSEM/CombineCounts.FilterLowExpression.tsv"
    
    sample_list_public_file <- "~/20220915_gastric_multiple/dna_combinePublic/public_ref/combine/MutationInfo.combine.tsv"
    raw_tcga_count_file <- "~/20220915_gastric_multiple/dna_combinePublic/public_ref/TCGA/TCGA-STAD.RSEM_counts.tsv"

    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/mRNA"
    gtf_file <- "~/ref/GTF/gencode.v19.ensg_genename.txt"
    gtf_for_len_file <- "~/ref/GTF/gencode.v19.annotation.gtf"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
rsem_file <- opt$rsem_file
sample_list_public_file <- opt$sample_list_public_file
raw_tcga_count_file <- opt$raw_tcga_count_file
out_path <- opt$out_path
gtf_file <- opt$gtf_file
gtf_for_len_file <- opt$gtf_for_len_file

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))
dat_tpm <- data.frame(fread(rsem_file))
dat_gtf <- data.frame(fread(gtf_file , header = F))
colnames(dat_gtf) <- c("gene_id" , "Hugo_Symbol")

info_public <- data.frame(fread(sample_list_public_file))
dat_tpm_public <- data.frame(fread(raw_tcga_count_file))

##########################################################################################

rownames(dat_tpm_public) <- dat_tpm_public$sample
dat_tpm_public <- dat_tpm_public[,-1]
## 转化为原始的counts

dat_tpm_public_raw <- 2^dat_tpm_public - 1
colnames(dat_tpm_public_raw) <- gsub( '[.]' , '-' , colnames(dat_tpm_public_raw))
## 提取最后纳入分析的样本
dat_tpm_public_raw <- dat_tpm_public_raw[,which(substring(colnames(dat_tpm_public_raw) , 1 , 12) %in% info_public$Tumor)]
colnames(dat_tpm_public_raw) <- gsub( '-11' , '_Normal' , colnames(dat_tpm_public_raw))

info_TCGA <- subset( info_public , From == "TCGA" )
info_TCGA$ID <- paste0(info_TCGA$Tumor , "-01")

for(id in unique(info_TCGA$ID)){
    class <- subset(info_TCGA , ID == id)$Class
    tumor <- subset(info_TCGA , ID == id)$Tumor
    if(length(class)!=0){
        colnames(dat_tpm_public_raw) <- gsub( id , paste0(tumor , "_" , class) , colnames(dat_tpm_public_raw))
    }
}

dat_tpm_public_raw_use <- dat_tpm_public_raw
#dat_tpm_public_raw_use <- dat_tpm_public_raw_use[,grep( "Normal" , colnames(dat_tpm_public_raw_use) , invert = T )]

##########################################################################################
## TCGA计算TPM，过滤低表达
## 计算基因长度
txdb <- makeTxDbFromGFF(gtf_for_len_file, format='gtf')
# then collect the exons per gene id
exons.list.per.gene <- exonsBy(txdb, by="gene")
# then for each gene, reduce all the exons to a set of non overlapping exons, 
# calculate their lengths (widths) and sum then
exonic.gene.sizes <- sum(width(reduce(exons.list.per.gene)))
gene_length <- data.frame(exonic.gene.sizes)
gene_length$ENSG <- substring( rownames(gene_length) , 1, 15 )
gene_length <- merge( gene_length , dat_gtf , by.x = "ENSG" , by.y = "gene_id" )
gene_length <- gene_length %>%
group_by( Hugo_Symbol ) %>%
summarize( length = max(exonic.gene.sizes) )

dat_tpm_public_raw_use_tpm <- dat_tpm_public_raw_use
dat_tpm_public_raw_use_tpm$Hugo_Symbol <- rownames(dat_tpm_public_raw_use_tpm)
dat_tpm_public_raw_use_tpm <- merge( gene_length , dat_tpm_public_raw_use_tpm )

## 计算TPM
Counts2TPM <- function(counts, effLen){
  rate <- log(counts) - log(effLen)
  denom <- log(sum(exp(rate)))
  exp(rate - denom + log(1e6))
}

result <- data.frame()
for(i in 3:ncol(dat_tpm_public_raw_use_tpm)){

    counts <- dat_tpm_public_raw_use_tpm[,i]
    effLen <- dat_tpm_public_raw_use_tpm[,"length"]
    tmp_tpm <- data.frame( Counts2TPM(counts, effLen))
    colnames(tmp_tpm) <- colnames(dat_tpm_public_raw_use_tpm)[i]
    rownames(tmp_tpm) <- dat_tpm_public_raw_use_tpm$Hugo_Symbol
    if(nrow(result) == 0 ){
        result <- tmp_tpm
    }else{
        result <- cbind(result , tmp_tpm)
    }
}

## 去除中位表达小于1的基因
keep_gene <- which(apply(result , 1 , median) > 1)
dat_tpm_public_raw_use_filter <- dat_tpm_public_raw_use[names(keep_gene) , ]

image_name <- paste0( out_path , "/TCGA.FilterLowExpression.TPM.tsv" )
write.table( result[keep_gene,] , image_name , row.names = T , sep = "\t" , quote = F )

##########################################################################################

class_type <- c( "Normal" , "IM" , "IGC" , "DGC")

##########################################################################################

dat_tpm_all <- dat_tpm[,-1]
colnames(dat_tpm_all) <- gsub( "[.]" , "-" ,  colnames(dat_tpm_all) )
rownames(dat_tpm_all) <- dat_tpm_all$Hugo_Symbol
dat_tpm_all <- dat_tpm_all[,-1]

##########################################################################################

share_gene <- rownames(dat_tpm_all)[rownames(dat_tpm_all) %in% rownames(dat_tpm_public_raw_use_filter)]
dat_combine <- cbind( dat_tpm_all[share_gene,] , dat_tpm_public_raw_use_filter[share_gene,] )

##########################################################################################
## 构建分组信息
condition <- sapply(strsplit(colnames(dat_combine) , "_") , "[" , 2)
condition <- factor(sapply(strsplit(condition , "-") , "[" , 1))    

#individual <- factor(sapply(strsplit(colnames(dat_tpm_all)[-1] , "_") , "[" , 1))
#coldata <- data.frame(row.names = colnames(dat_tpm_all)[-1] , condition , individual)

## counts变为整数
readCount <- round(dat_combine)

##########################################################################################
## 去除批次效应
## https://github.com/zhangyuqing/ComBat-seq
batch <- c(rep( "NJMU" , ncol(dat_tpm_all)) , rep( "TCGA" , ncol(dat_tpm_public_raw_use)))
adjusted_counts <- ComBat_seq(as.matrix(readCount), batch=batch, group=condition)

##########################################################################################
## Function
## 标化TMM
tmm_stand <- function(counts = counts , condition = condition) {
    ##########################################################################################
    ## https://rpubs.com/LiYumei/806213
    ## Count matrix preprocessing using edgeR package
    y <- DGEList(counts = counts , group = condition)

    ## 去除TPM低表达的基因
    ## Remove rows consistently have zero or very low counts
    #keep <- filterByExpr(y)
    #y <- y[keep,keep.lib.sizes=FALSE]

    ## Perform TMM normalization and transfer to CPM (Counts Per Million)
    ## 统一标化
    y <- calcNormFactors(y,method="TMM")
    count_norm <- cpm(y)
    count_norm <- as.data.frame(count_norm)

    ##########################################################################################
    ## 标准化的矩阵,用于后续样本间的表达比较
    ## 后续除软件要求计算评分使用TPM，否则均使用该矩阵
    normalized_counts <- count_norm
    normalized_counts <- data.frame(normalized_counts)
    normalized_counts$gene_id <- rownames(normalized_counts)
    colnames(normalized_counts) <- gsub( "[.]" , "-" ,  colnames(normalized_counts) )

    return(normalized_counts)
}

## 合并数据的标化
normalized_counts <- tmm_stand(counts = adjusted_counts , condition = condition)
image_name <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.tsv" )
write.table( normalized_counts , image_name , row.names = F , sep = "\t" , quote = F )

## TCGA数据单独标化
condition_tcga <- sapply(strsplit(colnames(dat_tpm_public_raw_use_filter) , "_") , "[" , 2)
normalized_tcga <- tmm_stand(counts = dat_tpm_public_raw_use_filter , condition = condition_tcga)
image_name <- paste0( out_path , "/TCGA.FilterLowExpression.TMM.tsv" )
write.table( normalized_tcga , image_name , row.names = F , sep = "\t" , quote = F )

##########################################################################################
## PCA的图
#################################################
## 检查批次效应是否去干净
pca.plot = function(dat , col , out_file){

    df.pca <- PCA(t(dat), graph = FALSE)
    fviz_pca_ind(df.pca,
       geom.ind = "point",
       col.ind = col ,
       addEllipses = TRUE,
       legend.title = "Groups"
    )
}

dat <- normalized_counts
disease <-  sapply(strsplit(sapply(strsplit(colnames(dat)[-ncol(dat)] , "_") , "[" , 2) , "-"  ) , "[" , 1)
Sample <- paste0( disease , "_" , batch )

## 现在的批次
out_file <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.PCA_Batch_Class.pdf" )
pdf(out_file)
pca.plot( dat[,-ncol(dat)], factor(Sample) , out_file)
dev.off()

## 现在的批次
out_file <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.PCA_Batch.pdf" )
pdf(out_file)
pca.plot( dat[,-ncol(dat)], factor(batch) , out_file)
dev.off()

id <- "Normal"
out_file <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.PCA_Batch_",id,".pdf" )
pdf(out_file)
tmp <- dat[,which(disease==id)]
from <- batch[which(disease==id)]
## 现在的批次
pca.plot( tmp, factor(from) , out_file)
dev.off()

id <- "IGC"
out_file <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.PCA_Batch_",id,".pdf" )
pdf(out_file)
tmp <- dat[,which(disease==id)]
from <- batch[which(disease==id)]
## 现在的批次
pca.plot( tmp, factor(from) , out_file)
dev.off()

id <- "DGC"
out_file <- paste0( out_path , "/CombineCounts.TCGA_NJMU.FilterLowExpression.TMM.PCA_Batch_",id,".pdf" )
pdf(out_file)
tmp <- dat[,which(disease==id)]
from <- batch[which(disease==id)]
## 现在的批次
pca.plot( tmp, factor(from) , out_file)
dev.off()
